import tensorflow as tf
from tensorflow.keras.datasets import mnist
import numpy as np
import matplotlib.pyplot as plt

# 加载 MNIST 数据集
_, (test_images, test_labels) = mnist.load_data()

# 调整数据形状以适应 LeNet5 的输入 (28*28 -> 32*32)
test_images = np.pad(test_images, ((0, 0), (2, 2), (2, 2)), 'constant', constant_values=0)
test_images = test_images.reshape((10000, 32, 32, 1))

# 归一化处理
test_images = test_images.astype('float32') / 255

# 加载保存的模型
model = tf.keras.models.load_model('./LeNet5/lenet5_model.h5')

# 进行推理
predictions = model.predict(test_images)

# 获取预测的类别
predicted_classes = np.argmax(predictions, axis=1)

# 随机选择一些样本进行可视化
num_samples = 9
random_indices = np.random.choice(test_images.shape[0], num_samples, replace=False)

plt.figure(figsize=(10, 10))
for i, index in enumerate(random_indices):
    plt.subplot(3, 3, i + 1)
    plt.imshow(test_images[index].reshape(32, 32), cmap='gray')
    plt.title(f"Pred: {predicted_classes[index]}, True: {test_labels[index]}")
    plt.axis('off')

plt.show()
